"""
model name : MY_YOLO
file       : detect0.py
information:
    author : OuYang
    time   : 2025/2/1
"""
import argparse
import os

import cv2
from PIL import Image

from model import YOLODetector

# Get current time
from datetime import datetime

import torch

current_time = datetime.now().strftime("%m-%d-%H-%M")


def parse_opt():
    parser = argparse.ArgumentParser()

    parser.add_argument('--dataset', type=str, default='D:/space/datasets/BITVehicle/valid/images',
                        help='path to dataset')
    parser.add_argument('--model', type=str, help='Model')
    parser.add_argument('--weights', type=str, help='weights', required=True)
    parser.add_argument('--num_classes', type=int, default=20, help='number of classes', required=True)
    parser.add_argument('--classes', type=str, default=None, help='class text', required=True)
    parser.add_argument('--imgsz', type=int, default=448, help='input image size')
    parser.add_argument('--backbone', type=str, default='resnet50', help='backbone')
    parser.add_argument('--s', type=int, default=7, help='s')
    parser.add_argument('--b', type=int, default=2, help='b')
    parser.add_argument('--save_dir', type=str, default=f'./runs/detect/{current_time}')
    parser.add_argument('--device', type=str, default='cuda:0', help='detect device type')
    parser.add_argument('--arg', type=str, default=None, help='argument text')
    parser.add_argument('--confidence', type=float, default=0.5, help='confidence')
    parser.add_argument('--scale', type=float, default=0.1, help='scale')
    parser.add_argument('--iou', type=float, default=0.5, help='iou')
    parser.add_argument('--show', action='store_true', default=True, help='show images')

    return parser.parse_args()


def detect(opt):
    # Save
    if not os.path.exists(opt.save_dir):
        os.makedirs(opt.save_dir)
        os.makedirs(os.path.join(opt.save_dir, "pictures"))
        os.makedirs(os.path.join(opt.save_dir, "videos"))

    # Print and Save info
    with open(os.path.join(opt.save_dir, "arguments.txt"), 'w') as file:
        for key, value in vars(opt).items():
            print(f'{key:20}=\t{value}')
            file.write(f"{key:20}=\t{value}\n")

    # Device
    device = torch.device(opt.device)
    if not torch.cuda.is_available():
        print("Error: Can't use GPU for detecting")
        device = torch.device('cpu')

    # Load class names
    class_names = []
    with open(opt.classes, 'r') as classes_file:
        for class_name in classes_file:
            class_name = class_name.strip()
            class_names.append(class_name)

        print(f"{'class_names':20}=\t{' '.join(class_names)}")

    # Load mod
    yolo = YOLODetector(
        weights=opt.weights,
        num_classes=opt.num_classes,
        classes_name=class_names,
        device=device,
        imgsz=opt.imgsz,
        confidence_threshold=opt.confidence,
        scale_threshold=opt.scale,
        iou_threshold=opt.iou,
        backbone=opt.backbone,
        s=opt.s,
        b=opt.b,
        show=opt.show,
    )

    # Detect
    for image_name in os.listdir(opt.dataset):
        # Read Image
        image_path = os.path.join(opt.dataset, image_name)
        image = Image.open(image_path).convert('RGB')

        # Detecting
        result_images = yolo.detect([image])

        # Save
        for result_image in result_images:
            cv2.imwrite(os.path.join(opt.save_dir, 'pictures', image_name), result_image)


if __name__ == '__main__':
    opt = parse_opt()
    detect(opt)
